# -*- coding=utf-8 -*-
# 2022.0114
# @zlk
import sys
import multiprocessing
import argparse
import os
import numpy as np
# add data read model
# add vcf data filter model


def read_vcf_data(infile,p1,p2,b1,b2):
    def get_geno(geno_str):
        num1=geno_str.split(':')[1].split(',')[0]
        num2=geno_str.split(':')[1].split(',')[1]
        return num1+'|'+num2
    out_dict={}
    vcf_file=open(infile,'r')
    out_table_file=open(infile+'.table','w')
    for line in vcf_file:
        if line.startswith('#'):
            continue
        line_list=line.strip().split()
        if ',' in line_list[3] or ',' in line_list[4]:
            continue
        # filter genotrype by parents
        if line_list[p1].split(':')[0]==line_list[p2].split(':')[0] or line_list[p1].startswith('./.') or line_list[p2].startswith('0/1') or line_list[p2].startswith('0|1'):
            continue
        p1_geno=get_geno(line_list[p1])
        p2_geno=get_geno(line_list[p2])
        b1_geno=get_geno(line_list[b1])
        b2_geno=get_geno(line_list[b2])
        out_table_file.write(line_list[0]+'\t'+line_list[1]+'\t'+line_list[3]+'\t'+line_list[4]+'\t'+p1_geno+'\t'+p2_geno+'\t'+b1_geno+'\t'+b2_geno+'\n')
        if line_list[0] not in out_dict.keys():
            out_dict[line_list[0]]=[]
        out_dict[line_list[0]].append([int(line_list[1]),line_list[3],line_list[4],p1_geno,p2_geno,b1_geno,b2_geno])
    vcf_file.close()
    return out_dict

def read_table_file(infile):
    table_file=open(infile,'r')
    out_dict={}
    for line in table_file:
        line_list=line.strip().split('\t')
        if line_list[0] not in out_dict.keys():
            out_dict[line_list[0]]=[]
        ad_list=[int(line_list[1])]+line_list[2:]
        out_dict[line_list[0]].append(ad_list)
    return out_dict

def pre_binom_test(rep):
    out_dict={}
    for n in range(5,500,1):
        num=np.percentile(np.random.binomial(n, 0.5, rep),99.95)
        out_dict[n]=num
    return out_dict

def filter_data(key,chr_list,dep_lim,dep_high,dep_lim_pool,dep_high_pool,binom_dict):
    out_snp_list=[]
    def get_dep(geno):
        # print(geno)
        dep=int(geno.split('|')[0])+int(geno.split('|')[1])
        geno=[int(geno.split('|')[0]),int(geno.split('|')[1])]
        return dep,geno
    def binom_index(b1_dep,b2_dep,binom_dict):
        num1=binom_dict[int(b1_dep)]
        index1=(num1-(b1_dep-num1))/b1_dep
        num2=binom_dict[int(b2_dep)]
        index2=(num2-(b2_dep-num2))/b2_dep
        return index1+index2
    for vcf_list in chr_list:
        p1_dep,p1_geno=get_dep(vcf_list[3])
        p2_dep,p2_geno=get_dep(vcf_list[4])
        b1_dep,b1_geno=get_dep(vcf_list[5])
        b2_dep,b2_geno=get_dep(vcf_list[6])
        # filter by depth
        if p1_dep<dep_lim or p1_dep>dep_high or p2_dep>dep_high or b1_dep<dep_lim_pool or b1_dep>dep_high_pool or b2_dep<dep_lim_pool or b2_dep>dep_high_pool:
            continue
        # filter by p1 genotype
        if p1_geno[0]>binom_dict[p1_dep] or p1_geno[1]>binom_dict[p1_dep]:
            continue
        # fileter bu P2 genotype
        if p2_dep<2:
            flag='PAV'
            b1_index=(b1_geno[0]-b1_geno[1])/b1_dep
            b2_index=(b2_geno[0]-b2_geno[1])/b2_dep
            total_index=abs(b1_index-b2_index)
        elif p2_geno[0]==0 or p2_geno[1]==0  or p2_geno[1]/p2_geno[0]<0.1 or p2_geno[1]/p2_geno[0]>10:
            if p2_geno[1]>p2_geno[0]:
                p1_num,p2_num=0,1
            else:
                p1_num,p2_num=1,0
            if len(vcf_list[1])>1 or  len(vcf_list[2])>1:
                flag='indel'
            else:
                flag='snp'
            b1_index=(b1_geno[p2_num]-(b1_dep/2)-b1_geno[p1_num])/(b1_dep/2)
            b2_index=(b2_geno[p2_num]-(b2_dep/2)-b2_geno[p1_num])/(b2_dep/2)
            if abs(b1_index)>1.1 or abs(b2_index)>1.1:
                continue
            total_index=abs(b1_index-b2_index)
        else:
            continue
        #####
        virtual_index=binom_index(b1_dep/2,b2_dep/2,binom_dict)
        out_snp_list.append(vcf_list+[virtual_index,total_index,flag])
    return [key,out_snp_list]

def read_dhhp_file(infile):
    loop_file=open(infile,'r')
    out_dict={}
    for line in loop_file:
        line_list=line.strip().split('\t')
        if line_list[0] not in out_dict.keys():
            out_dict[line_list[0]]=[]
        ad_list=[int(line_list[1])]+line_list[2:]
        out_dict[line_list[0]].append(ad_list)
    return out_dict

def cal_dis_func(key,win_size,loop_value):
    step=win_size/10
    win_list=[0,win_size]
    total_list=[]
    start_pos=0
    # this_list=[]
    while win_list[1]<loop_value[-1][0]:
        this_list=[[],[]]
        for loop_list in loop_value[start_pos:]:
            # print(loop_list)
            i_list=[int(loop_list[0]),float(loop_list[-3]),float(loop_list[-2])]
            if win_list[0]<=i_list[0]<win_list[1]:
                this_list[0].append(i_list[1])
                this_list[1].append(i_list[2])
                # this_list[2].append(i_list[3])
            elif i_list[0]>=win_list[1]:
                try:
                    total_list.append([(win_list[0]+win_list[1])/2,sum(this_list[0])/len(this_list[0]),sum(this_list[1])/len(this_list[0]),\
                        len(this_list[0])]) 
                except ZeroDivisionError:
                    total_list.append([(win_list[0]+win_list[1])/2,0,0,0])

                win_list[0]+=step
                win_list[1]+=step
                break
            elif i_list[0]<win_list[0]:
                start_pos+=1
    return [key,total_list]


if __name__=="__main__":
    parser = argparse.ArgumentParser(description='输入参数如下:')
    parser.add_argument('-p1', '--parent1', type=int, help='显性亲本在vcf中的列数',)
    parser.add_argument('-p2', '--parent2', type=int, help='隐性亲本在vcf中的列数',)
    parser.add_argument('-b1', '--pool1', type=int, help='与显性亲本相同表型的混池在vcf中的列数',)
    parser.add_argument('-b2', '--pool2', type=int, help='与隐性亲本相同表型的混池在vcf中的列数',)
    parser.add_argument('-d1', '--parentdep1', type=int, help='亲本的最低覆盖度', default=10)
    parser.add_argument('-d2', '--pooldep1', type=int, help='混池的最低覆盖度', default=20)
    parser.add_argument('-d3', '--parentdep2', type=int, help='亲本的最高覆盖度', default=80)
    parser.add_argument('-d4', '--pooldep2', type=int, help='混池的最高覆盖度', default=200)
    parser.add_argument('-w', '--win', type=int, help='滑窗的大小',default=1000000)
    parser.add_argument('-vcf', '--input_vcf', help='vcf文件 必要参数')
    parser.add_argument('-table', '--input_format', help='本程序简易的VCF文件格式, 替换VCF文件的输入')
    parser.add_argument('-dhhp', '--input_DHHP', help='DHHP文件, 如果只是调整窗口大小, 可以使用之前生成的中间文件(.DHHP)')
    parser.add_argument('-o', '--out', help='输出文件名', required=True)
    args = parser.parse_args()
    p1=args.parent1
    p2=args.parent2
    b1=args.pool1
    b2=args.pool2
    dep_lim=args.parentdep1
    dep_high=args.parentdep2
    dep_lim_pool=args.pooldep1
    dep_high_pool=args.pooldep2
    win_size=args.win
    outfile=args.out    
    infile_vcf=args.input_vcf
    infile_bsa=args.input_DHHP
    infile_format=args.input_format
    if infile_bsa==None:
        if p1==None or p2==None or b1==None or b2==None:
            print('You have to provide -p1 -p2 -b1 -b2')
    if infile_vcf==None and infile_bsa==None and infile_format==None:
        print('You have to provide .vcf or .bsa file')
        sys.exit()
    ###get binom test dict
    binom_dict=pre_binom_test(1000)
    if infile_bsa:
        print('Start to read the file '+ infile_bsa)
        chr_geno_dict_filter=read_dhhp_file(infile_bsa)
    else:
    # chr_geno_dict={}
        if infile_vcf:
            print('Start reading file '+infile_vcf)
            chr_geno_dict=read_vcf_data(infile_vcf,p1,p2,b1,b2)
        elif infile_format:
            print('Start reading file '+infile_format)
            chr_geno_dict=read_table_file(infile_format)
        chr_geno_dict_filter={}
        ##########filter data
        print('Start calculating the DHHP-value')
        pool = multiprocessing.Pool(len(chr_geno_dict))
        results=[]
        for key,value in chr_geno_dict.items():

            print(key)
            # filter_list=filter_data(key,value,dep_lim,dep_high,dep_lim_pool,dep_high_pool,binom_dict)
            # chr_geno_dict_filter[key]=filter_list
            results.append(pool.apply_async(filter_data, (key,value,dep_lim,dep_high,dep_lim_pool,dep_high_pool,binom_dict)))
        pool.close()
        pool.join()
        for result in results:
            chr_list=result.get()
            ########## filter by snp number ###########
            if len(chr_list[1])<10000:
                continue
            chr_geno_dict_filter[chr_list[0]]=chr_list[1]
        ####bsa file output
        out_bsa_file=open(outfile+'.dhhp','w')
        for key,value in chr_geno_dict_filter.items():
            for i_list in value:
                out_bsa_file.write(key+'\t')
                for i in i_list:
                    out_bsa_file.write(str(i)+'\t')
                out_bsa_file.write('\n')
        out_bsa_file.close()
        print('The calculation has been completed and the results are stored in file '+outfile+'.dhhp')
    
    #######window####
    print('Start smoothing dhhp value, window:'+str(win_size))
    pool = multiprocessing.Pool(len(chr_geno_dict_filter))
    result=[]
    for key,value in chr_geno_dict_filter.items():
        print(key)
        # cal_dis_func(key,win_size,value)
        result.append(pool.apply_async(cal_dis_func,(key,win_size,value)))
    pool.close()
    pool.join()
    file3=open(outfile,'w')
    for i_list in result:
        i_list=i_list.get()
        for i in i_list[1][1:]:
            file3.write(i_list[0]+'\t'+str(int(i[0]))+'\t')
            for x in i[1:]:
                file3.write(str(x)+'\t')
            file3.write('\n')
    print('All jobs have been completed and the smoothed results are stored in '+outfile)


